# %%
import numpy as np  # 导入整个numpy模块并用np作为它的新名字，要使用句点表示法访问需要的类
import matplotlib.pyplot as plt  # matplotlib.pyplot是matplotlib的子库，它提供了一种类似于MATLAB的绘图方式
from matplotlib.font_manager import FontProperties  # 用于跨平台查找、管理和使用字体的模块
from math import *  # 导入math模块的每个类，可以直接使用类，无需句点表示法，后边就可以直接用sqrt(), pi, exp()等
import cmath  # 导入cmath模块用来处理复数情况，如cmath.exp()


# 定义哈密顿量
def hamiltonian(kx, ky, epsilon=-1.01354, t=2.01493, r=-0.337845):  # 石墨烯的TB哈密顿量，epsilon, t, r不赋值将默认使用括号中的值
    h0 = np.zeros((2, 2)) * (1 + 0j)  # 产生复数格式的2*2的矩阵
    h1 = np.zeros((2, 2)) * (1 + 0j)
    h2 = np.zeros((2, 2)) * (1 + 0j)

    # on-site能
    h0[0, 0] = epsilon
    h0[1, 1] = epsilon

    # 最近邻
    h1[1, 0] = t * (cmath.exp(1j * ky / sqrt(3)) + cmath.exp(1j * (-kx / 2 - ky / 2 / sqrt(3))) + cmath.exp(
        1j * (kx / 2 - ky / 2 / sqrt(3))))
    h1[0, 1] = h1[1, 0].conj()

    # 次近邻
    h2[0, 0] = 2 * r * cos(kx) + 4 * r * cos(kx / 2) * cos(sqrt(3) * ky / 2)
    h2[1, 1] = 2 * r * cos(kx) + 4 * r * cos(kx / 2) * cos(sqrt(3) * ky / 2)

    h = h0 + h1 + h2
    return h


# %%
# 数值计算
Gamma = np.array([0, 0])
M = np.array([pi, pi / sqrt(3)])
K = np.array([2 * pi / 3, 2 * sqrt(3) * pi / 3])

kn = 100  # 每个区域的取点数，三个区域分别为Gamma-M, M-K, K-Gamma
n = 3  # n个区域
kx_array = np.zeros(kn * n)  # kx对应总共有kn*n个点
ky_array = np.zeros(kn * n)  # ky对应总共有kn*n个点

# Gamma-M
kx_array[0:kn] = np.linspace(0, M[0], kn)
ky_array[0:kn] = np.linspace(0, M[1], kn)

# M-K
kx_array[kn:2 * kn] = np.linspace(M[0], K[0], kn)
ky_array[kn:2 * kn] = np.linspace(M[1], K[1], kn)

# K-Gamma
kx_array[2 * kn:3 * kn] = np.linspace(K[0], Gamma[0], kn)
ky_array[2 * kn:3 * kn] = np.linspace(K[1], Gamma[1], kn)

i0 = 0
dim = hamiltonian(0, 0).shape[0]  # hamiltonian(0,0).shape得到(2,2)，然后hamiltonian(0,0).shape[0]得到2
eigenvalue_k = np.zeros((n * kn, dim))  # 哈密顿矩阵是2维矩阵，它有两个本征值，路径有300个k点，所以这里总的本征值维度为(300,2)
fig, ax = plt.subplots()  # fig为返回的图像，ax为返回的坐标系（为一个数组）
for kn0 in range(n * kn):
    kx = kx_array[kn0]
    ky = ky_array[kn0]
    eigenvalue, eigenvector = np.linalg.eig(hamiltonian(kx, ky))  # 计算矩阵的本征值和本征向量
    eigenvalue_k[i0, :] = np.sort(np.real(eigenvalue[:]))  # 对给定的数组元素进行排序
    # 这里对本征值的实部排序，本征值应该是实数，这里用real是为了防止很小的虚数干扰，因为Python进行进制转换会导致精度误差
    i0 += 1

# 画图
# k路径
kpath = np.zeros(kn * n)
kpath[0:kn] = np.linspace(0, np.linalg.norm(M), kn)
kpath[kn:2 * kn] = np.linspace(kpath[kn - 1], kpath[kn - 1] + np.linalg.norm(K - M), kn)
kpath[2 * kn:3 * kn] = np.linspace(kpath[2 * kn - 1], kpath[2 * kn - 1] + np.linalg.norm(Gamma - K), kn)

for dim0 in range(dim):  # 本征值有两组，本征值组数等于矩阵维度
    color = ['-r', '-b']  # 设置能带线条样式，能带线条样式的数目等于本征值组数，'-r'表示红色实线，'-b'表示蓝色实线
    plt.plot(kpath, eigenvalue_k[:, dim0], color[dim0])

# 画图时中文标题无法显示，需要添加中文路径
# plt.title('The band structure of Graphene')
# plt.title('石墨烯能带', fontproperties=FontProperties(fname='/System/Library/Fonts/Supplemental/Songti.ttc'))  # Mac 宋体
# plt.title('石墨烯能带', fontproperties=FontProperties(fname='/System/Library/Fonts/PingFang.ttc'))  # Mac 苹方-简
# plt.title('石墨烯能带', fontproperties=FontProperties(fname='/System/Library/Fonts/STHeiti Light.ttc'))   # Mac 黑体-简
plt.title('石墨烯能带', fontproperties=FontProperties(fname='\\WINDOWS\\FONTS\\SIMSUN.TTC'))  # Windows 宋体

plt.ylabel("Energy")
# plt.ylabel("能量", fontproperties=FontProperties(fname='/System/Library/Fonts/Supplemental/Songti.ttc'))
ax.set_xticks([0, kpath[kn - 1], kpath[2 * kn - 1], kpath[3 * kn - 1]])  # 设置x轴刻度
ax.set_xticklabels([r'$\Gamma$', 'M', 'K', r'$\Gamma$'])  # 设置x轴坐标刻度
plt.xlim(0, kpath[3 * kn - 1])  # 坐标轴范围
plt.grid(axis="x", c="c", linestyle="--")  # 生成网格，以x轴坐标生成，线条为红色短线
plt.savefig('GrapheneBand.png', dpi=600)
plt.show()
